# PDFDecryper-4.0
# Auther:tianwenxiao
# mail:tianwenxiao@wo.cn
# 主要功能:
#        2.0：针对没有设置打开密码的PDF文件，只是设置了打印、编辑等限制的密码，可以使用该脚本进行解密
#        3.5: 新增：针对设置了打开密码的文件，增加了密码字典库爆破功能，只需要选择需要破解的PDF文件，再选择用于破解的字典库，点击执行脚本即可开始破解
#        4.0：新增：针对需要使用密码字典库进行爆破的操作，新增了多线程的方式，通过生成器将字典文件分成块，使用多线程分别遍历，极大的提高了破解效率，但是
#            目前仍然有个问题，那就是目前GUI上的进度条目前无法更新，后续版本即将修复。
from shlex import join
import threading
import tkinter as tk
from tkinter import filedialog
from tkinter import messagebox
from uu import Error
#pdf的读取方法
from PyPDF2 import PdfReader
#pdf的写入方法
from PyPDF2 import PdfWriter
#高加密的方法
from Crypto.Cipher import AES
from tkinter import ttk
import concurrent.futures
import queue

# 定义一个类在线程中共享信息
class SharedState:
    def __init__(self):
        self.state = True
        self.lock = threading.Lock()

    def set_state(self, new_state):
        with self.lock:
            self.state = new_state

    def get_state(self):
        with self.lock:
            return self.state

# 创建一个线程间信息通信对象
thread_staus = SharedState()

# 切分密码库文件的生成器
# 生成器
def chunks_generator(filename, chunk_size):
    with open(filename, 'r', encoding='utf-8') as file:
        chunk = []
        for line in file:
            line = line.strip()
            chunk.append(line)
            if len(chunk) == chunk_size:  # 修改这里
                yield chunk
                chunk = []  # 清空chunk以便收集下一个数据块
        if chunk:  # 处理文件末尾不完整的块
            yield chunk

# 更新进度条的方法
def check_progress(root,progress, progress_label, progress_queue, total_tasks):
    current_state = thread_staus.get_state()
    if current_state:
        if not progress_queue.empty():
            # 阻塞调用，如果队列为空，则等待
            completed_work = progress_queue.get(block=True)
            current_value = progress['value'] + completed_work
            progress['value'] = current_value
            progress_label.config(text=f"Progress: {current_value / total_tasks * 100:.1f}%")
            print("刷新进度条{}".format(current_value))
            if current_value > total_tasks:
                current_state = False
                thread_staus.set_state(False)
            root.update_idletasks()
    else:
        root.update_idletasks()

# 统计密码库总共有多少行
def count_lines_in_file(file_path):
    with open(file_path, 'r', encoding='utf-8') as file:
        return sum(1 for _ in file)

# 选择要处理的文件
def select_input_file(input_file_path_var):
    input_file_path = filedialog.askopenfilename(filetypes=[("PDF files", "*.PDF;*.pdf")])
    if input_file_path:
        input_file_path_var.delete(0, tk.END)
        input_file_path_var.insert(0, input_file_path)
    return input_file_path

# 选择用于破解的密码字典
def select_cracker_dic(input_crack_dic_var):
    crack_dic = filedialog.askopenfilename(filetypes=[("TXT file","*.txt;*.TXT")])
    if crack_dic:
        input_crack_dic_var.delete(0,tk.END)
        input_crack_dic_var.insert(0,crack_dic)
        return crack_dic

# 定义一个函数来读取PDF文件
def get_reader(filename,password=None):
    try:
        old_file = open(filename, 'rb')
    except Exception as err:
        return print('文件打开失败！' + str(err))
    #如果使用python2需要将PdfReader改为PdfFileReader
    pdf_reader = PdfReader(old_file, strict=False)

    # 如果使用python2需要将is_encrypted改为isEncrypted
    # 执行解密操作
    if pdf_reader.is_encrypted:
        if password is None:
            print('文件被加密，需要密码！--{}'.format(filename))
            return pdf_reader
        else:
            success = pdf_reader.decrypt(password)
            if not success :
                #return print('密码不正确！--{}'.format(filename))
                return
            else:
                #print("解密成功，密码为{}".format(password))
                return pdf_reader
    elif old_file in locals():
        old_file.close()
        # 返回结果
    return pdf_reader

# 真正处理PDF的方法
def deception_pdf(root,filename,progress_queue=None,count_lines=None,chunk=None) :
    current_state = thread_staus.get_state()
    if current_state:
        # 如果用户没有选择破解字典库
        if chunk is None or chunk == '' or count_lines is None or count_lines == '' or progress_queue == None:
            print("==没有选择破解字典库==")
            try:
                pdf_reader = get_reader(filename)
                if pdf_reader is None:
                    return print("无内容读取")
                # 如果使用python2需要将is_encrypted改为isEncrypted
                elif not pdf_reader.is_encrypted:
                    return print('文件没有被加密，无需操作')
                # 如果使用的是python2需要将PdfWriter改为PdfFileWriter
                pdf_writer = PdfWriter()
                #如果使用的是python2需要将将append_pages_from_reader改为appendPagesFromReader
                pdf_writer.append_pages_from_reader(pdf_reader)
                #创建解密后的pdf文件和展示文件的路径
                decrypted_filename = "".join(filename.split('.')[:-1]) + '_' + '已解密' + '.pdf'
                print("解密文件已生成:{}".format(decrypted_filename))
                # 写入新文件
                pdf_writer.write(open(decrypted_filename, 'wb'))
                current_state = False
                thread_staus.set_state(False)
                root.destroy()
            except Exception as err:
                print("文件设置了打开密码，无法读取PDF，请尝试选择字典库进行爆破")
                current_state = False
                thread_staus.set_state(False)
                #messagebox.showinfo("解密失败", "文件设置了打开密码，无法读取PDF，请尝试选择字典库进行爆破")
                root.destroy()
        else :
            print("选择了破解字典库==>{}".format(count_lines))
            count = 0
            for password in chunk:
                    if current_state:
                        count = count + 1
                        if count % 50 == 0:
                            # 同步信息至队列
                            progress_queue.put(count)
                        pdf_reader = get_reader(filename, password)
                        # 如果破解成功
                        if(pdf_reader):
                            # 如果使用的是python2需要将PdfWriter改为PdfFileWriter
                            pdf_writer = PdfWriter()
                            #如果使用的是python2需要将将append_pages_from_reader改为appendPagesFromReader
                            pdf_writer.append_pages_from_reader(pdf_reader)
                            #创建解密后的pdf文件和展示文件的路径
                            decrypted_filename = "".join(filename.split('.')[:-1]) + '_' + '已解密' + '.pdf'
                            # 写入新文件
                            pdf_writer.write(open(decrypted_filename, 'wb'))
                            thread_staus.set_state(False)
                            current_state = False
                            #messagebox.showinfo("解密完成", "您的PDF文件已破解完成，密码为{}".format(password))
                            break
                        else:
                            if(count<count_lines):
                                current_thread = threading.current_thread()  # 获取当前线程实例
                                thread_name = current_thread.getName()       # 获取当前线程名字
                                print(("线程{}正在解析密码==>{}".format(thread_name,password)))
                                pass
                            else:
                                current_state = False
                                thread_staus.set_state(False)
                                #messagebox.showinfo("标题", "抱歉,未找到匹配的密码，请换个字典库再试一下吧")
                                print("抱歉,未找到匹配的密码，请换个字典库再试一下吧")
                                break
    else:
        root.destroy()

# 取消破解，退出应用
def closeApp(root):
    thread_staus.set_state(False)
    root.destroy()

# 开启线程，开始破解
def start_deception_pdf_thread(root, filename, progress, progress_label,crack_dic=None):
    # 创建一个线程安全的队列来传递进度信息
    progress_queue = queue.Queue()
    # 重置终止标志
    thread_staus.set_state(True)
    print(crack_dic)

    if crack_dic == None or crack_dic == "":
        deception_pdf(root,filename, progress,progress_label)
    else:
        # 获取字典库总行数，并传到函数里面，用于显示进度条
        count_lines = count_lines_in_file(crack_dic)
        chunks = chunks_generator(crack_dic,1000)
        # 使用 ThreadPoolExecutor 管理线程池
        with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
            # 将生成器的每一行作为任务提交到线程池
            futures = [executor.submit(deception_pdf,root,filename, progress_queue,count_lines,chunk)for chunk in chunks]
            # 等待所有任务完成（可选）使用这个会导致在遍历完后，程序卡死
            #concurrent.futures.wait(futures)
            # 确保所有进度更新都已完成
            progress_queue.join()
            root.after(20,lambda:check_progress(root,progress, progress_label, progress_queue,count_lines))

# 创建主函数
def main():
    # 创建主窗口
    root = tk.Tk()
    root.title("PDF解密小工具")
    root.geometry("600x300")
    root.resizable(False, False)
    # 创建一个Tkinter变量来存储文件路径
    #input_file_path_var = tk.StringVar()
    # 输入文件框
    input_frame = tk.Frame(root)
    input_frame.grid(column=0, row=0, padx=0, pady=10,sticky="w")
    input_label = tk.Label(input_frame, text="输入文件:")
    input_label.grid(column=0, row=0, padx=17, pady=10,sticky="w")
    input_file_path_var = tk.Entry(input_frame, width=50)
    input_file_path_var.grid(column=1, row=0, padx=10, pady=10,sticky="e")
    input_button = tk.Button(input_frame, text="选择文件", command=lambda:select_input_file(input_file_path_var))
    input_button.grid(column=2, row=0, padx=6, pady=10,sticky="e")
    # 选择密码字典
    dic_frame = tk.Frame(root)
    dic_frame.grid(column=0, row=1, padx=0, pady=10,sticky="w")
    dic_frame_label = tk.Label(dic_frame, text="请选择密码字典:")
    dic_frame_label.grid(column=1, row=1, padx=0, pady=10,sticky="w")
    input_crack_dic_var = tk.Entry(dic_frame, width=50)
    input_crack_dic_var.grid(column=2, row=1, padx=10, pady=10,sticky="e")
    dic_frame_button = tk.Button(dic_frame, text="选择文件", command=lambda:select_cracker_dic(input_crack_dic_var))
    dic_frame_button.grid(column=3, row=1, padx=5, pady=10,sticky="e")
    # 创建进度条
    progress_frame= tk.Frame(root)
    progress_frame.grid(column=0, row=4, padx=5, pady=10,sticky="w")
    progress = ttk.Progressbar(progress_frame, style='custom.Horizontal.TProgressbar', orient='horizontal', length=450, mode='determinate')

    # 创建并配置进度条样式
    style = ttk.Style()
    # 定义进度条的样式（这里以设置背景色为例）
    style.configure('custom.Horizontal.TProgressbar', background='blue')
    progress.grid(column=0, row=4, padx=0, pady=10)
    progress_label = tk.Label(progress_frame, text="已完成0%")
    progress_label.grid(column=1, row=4,padx=10, pady=10,sticky="w")
    controller_frame=tk.Frame(root)
    controller_frame.grid(column=0, row=3, padx=0, pady=10)

    # 创建按钮用于执行脚本
    execute_button = tk.Button(controller_frame, text="执行脚本", command=lambda:start_deception_pdf_thread(root,input_file_path_var.get(),progress,progress_label,input_crack_dic_var.get()))
    execute_button.grid(column=0, row=3, padx=10, pady=10)
    exit_button = tk.Button(controller_frame,text="取消并退出",command=lambda:closeApp(root))
    exit_button.grid(column=1, row=3, padx=10, pady=10)
    # 运行Tkinter事件循环
    root.mainloop()

if __name__ == "__main__":
    main()
